# Make graphs for figure 1

# Load Models and Store AUCs
countries <- c('indo', 'colo')
dvs <- c("any", "high", "spike")

for (country in countries) {
  for (v in dvs) {
    source(paste("config_code/",country,"_data_setup.R",sep = ""))
    source(paste("config_code/",country,"_predictor_vars.R",sep = ""))
    models <- names(figures_1_2)[-1]
    T <- end.year - start.year + 1
    
    # Get population benchmark dataset
    pop_auc_array <- array(NA,
                           dim = T)
    filename <- paste(modeldir,"/",
                      country,
                      "_ebma_",
                      v,
                      "_population_onset.RData",
                      sep = "")
    load(filename)
    for (i in 1:T) {
      predictions <- as.vector(ebma.results[[i]]$fit.oos)
      realizations <- as.vector(ebma.results[[i]]$actual.oos)
      pop_auc_array[i] <- roc(response = realizations,
                              predictor = predictions)[['auc']]
    }
    
    auc_array <- array(NA,
                       dim = c(length(models),
                               T),
                       dimnames = list(models,
                                       NULL))
    for (x in models) {
      filename <- paste(modeldir,"/",
                        country,
                        "_ebma_",
                        v,
                        "_",x,"_onset.RData",
                        sep = "")
      load(filename)
      for (i in 1:T) {
        predictions <- as.vector(ebma.results[[i]]$fit.oos)
        realizations <- as.vector(ebma.results[[i]]$actual.oos)
        auc_array[x, i] <- (roc(response = realizations,
                                predictor = predictions)[['auc']]
                            - pop_auc_array[i])
      }
    }
    mean_auc <- apply(auc_array, 1, mean)
    
    # Make Graph
    var.order1 <- order(mean_auc[names(figures_1_2_fixed)[-1]],
                        decreasing = FALSE)
    var.order2 <- order(mean_auc[names(figures_1_2_slow)],
                        decreasing = FALSE)
    var.order3 <- order(mean_auc[names(figures_1_2_annual)],
                        decreasing = FALSE)
    ranks1 = rep(NA,length(var.order1))
    ranks1[var.order1] <- ((1:length(var.order1))
                           +length(var.order2)
                           +length(var.order3)
                           +ifelse(length(var.order2)>0,
                                   2,1))
    ranks2 = rep(NA,length(var.order2))
    ranks2[var.order2] <- 1:length(var.order2)+length(var.order3)+1
    ranks3 = rep(NA,length(var.order3))
    ranks3[var.order3] <- 1:length(var.order3)
    ranks <- c(ranks1,ranks2,ranks3)
    varsets.total <- c(names(figures_1_2_fixed)[-1],
                       names(figures_1_2_slow),
                       names(figures_1_2_annual))
    pdf(paste("graphs/",
              country,
              "_AUCs_",
              v,
              "_ebma_popplus_onset.pdf",
              sep=""),
        width = 7, 
        height = 7)
    par(mar=c(5, 12, 2, 1) + 0.1)
    plot(mean_auc[varsets.total],
         ranks,
         xlim = c(-.1,.1),
         ylim = c(0,max(ranks)+ifelse(length(var.order2)>0,
                                      2,1)),
         ylab = "",
         xlab = "Performance improvement over population (AUC)",
         yaxt='n',
         main = "")
    for (i in 1:length(ranks)) {
      lines(x = c(-1,1),
            y = c(ranks[i],ranks[i]),
            col = "gray95")
    }
    points(as.vector(t(auc_array[varsets.total,])),
           rep(ranks,each = T),
           pch = rep(c(17,
                       rep(16,T-2),
                       15),
                     length(varsets.total)),
           cex = rep(c(1,
                       rep(.5,T-2),
                       1),
                     length(varsets.total)),
           col = rep(c("gray",
                       rep("black",T-2),
                       "gray"),
                     length(varsets.total)))
    points(mean_auc[varsets.total],
           ranks)
    axis(2,
         at = ranks,
         labels = labels[varsets.total],
         las = 2)
    axis(2,
         at = c(length(ranks)+ifelse(length(var.order2)>0,
                                     3,2),
                length(ranks3)+1,
                length(ranks2)+length(ranks3)+ifelse(length(ranks2)>0,2,1)
         ),
         labels = c("Time-Invariant",
                    "Time varying (Annual)",
                    ifelse(length(ranks2)>0,"Slow-Moving","")),
         las = 2,
         font=4 )
    lines(x = c(0,0),
          y = c(-1,length(ranks)+ifelse(length(ranks2)>0,5,3)),
          lty = 3)
    dev.off()
  }
}
